#include "nn_detector.hpp"

#include <NvOnnxParser.h>

#include <algorithm>
#include <fstream>
#include <map>
#include <stdexcept>
#include <vector>

#include "cuda_runtime_api.h"
#include "nn_yolo.hpp"
#include "opencv2/opencv.hpp"
#include "yololayer.h"

float Sigmoid(float x) { return 1 / (1 + std::exp(-x)); }

static float IOU(const TRT::Detection &det1, const TRT::Detection &det2) {
  const float left =
      std::max(det1.x_ctr - det1.w / 2., det2.x_ctr - det2.w / 2.);
  const float right =
      std::min(det1.x_ctr + det1.w / 2., det2.x_ctr + det2.w / 2.);
  const float top =
      std::max(det1.y_ctr - det1.h / 2., det2.y_ctr - det2.h / 2.);
  const float bottom =
      std::min(det1.y_ctr + det1.h / 2., det2.y_ctr + det2.h / 2.);

  if (top > bottom || left > right) return 0.;

  const float inter_box_s = (right - left) * (bottom - top);
  const float ratio =
      inter_box_s / (det1.w * det1.h + det2.w * det2.h - inter_box_s);
  // SPDLOG_INFO("Ratio : {}", ratio);
  return ratio;
}

void NonMaxSuppression(std::vector<TRT::Detection> &dets, float nms_thresh) {
  if (dets.empty()) return;

  std::vector<TRT::Detection> keep;

  std::sort(dets.begin(), dets.end(),
            [](const TRT::Detection &det1, const TRT::Detection &det2) {
              return det1.conf > det2.conf;
            });

  for (std::size_t m = 0; m < dets.size(); ++m) {
    auto &highest = dets[m];
    keep.emplace_back(highest);
    for (size_t n = m + 1; n < dets.size(); ++n) {
      if (IOU(highest, dets[n]) > nms_thresh) {
        dets.erase(dets.begin() + n);
        --n;
      }
    }
  }
  dets = keep;
  SPDLOG_WARN("size : {}", dets.size());
}

void TRTDetector::PreProcess(cv::Mat img) {
  if (img.size() != cv::Size(Yolo::INPUT_H, Yolo::INPUT_W))
    cv::resize(img, img, cv::Size(Yolo::INPUT_H, Yolo::INPUT_W));

  for (std::size_t c = 0; c < 3; c++) {
    for (std::size_t h = 0; h < Yolo::INPUT_H; h++) {
      for (std::size_t w = 0; w < Yolo::INPUT_W; w++) {
        data_[(2 - c) * Yolo::INPUT_W * Yolo::INPUT_H + h * Yolo::INPUT_W + w] =
            (float)img.at<cv::Vec3b>(h, w)[c] / 255.0;
      }
    }
  }

  SPDLOG_DEBUG("PreProcessed");
}

std::vector<TRT::Detection> TRTDetector::PostProcess(std::vector<float> prob,
                                                     int width, int height) {
  int det_size = sizeof(TRT::Detection) / sizeof(float);

  std::vector<TRT::Detection> dets;
  float ratio_width = Yolo::INPUT_W / (width * 1.0f);
  float ratio_height = Yolo::INPUT_H / (height * 1.0f);
  for (int i = 0; i < prob[0] && i < Yolo::MAX_OUTPUT_BBOX_COUNT; i++) {
    if (prob[1 + det_size * i + 4] > conf_thresh_) {
      TRT::Detection det{
          .x_ctr = prob[1 + det_size * i + 0] / ratio_width,
          .y_ctr = prob[1 + det_size * i + 1] / ratio_height,
          .w = prob[1 + det_size * i + 2] / ratio_width,
          .h = prob[1 + det_size * i + 3] / ratio_height,
          .conf = prob[1 + det_size * i + 4],
          .class_id = prob[1 + det_size * i + 5],
      };
      dets.emplace_back(det);
    }
  }

  SPDLOG_WARN("PostProcessed {} < {}", dets.size(), prob[0]);

  // [4] obj conf
  // [5] [6] [7] [8] class_conf
  // [5] [6] [7] [8] *= [4] conf = obj_conf * class_conf;
  // [0] [1] [2] [3] = x_ctr y_ctr w h
  return dets;
}

bool TRTDetector::CreateContex() {
  SPDLOG_DEBUG("[TRTDetector] CreateContex.");
  context_ = UniquePtr<IExecutionContext>(engine_->createExecutionContext());
  if (!context_) {
    SPDLOG_ERROR("[TRTDetector] CreateContex Fail.");
    return false;
  }
  SPDLOG_DEBUG("[TRTDetector] CreateContex OK.");
  return true;
}

bool TRTDetector::InitMemory() {
  idx_in_ = engine_->getBindingIndex(Yolo::INPUT_BLOB_NAME);
  idx_out_ = engine_->getBindingIndex(Yolo::OUTPUT_BLOB_NAME);

  auto dim_in = engine_->getBindingDimensions(idx_in_);
  auto dim_out = engine_->getBindingDimensions(idx_out_);
  SPDLOG_INFO("idx in: {} out : {}", idx_in_, idx_out_);

  data_ = new float[Yolo::BATCH_SIZE * 3 * 640 * 640];

  SPDLOG_WARN("DATA : {}", sizeof(data_));
  for (int i = 0; i <= dim_in.MAX_DIMS; i++)
    SPDLOG_INFO("[{}] in: {} out : {}", i, dim_in.d[i], dim_out.d[i]);
  SPDLOG_INFO("[nbDims] in: {} out : {}", dim_in.nbDims, dim_out.nbDims);

  for (int i = 0; i < engine_->getNbBindings(); ++i) {
    Dims dim = engine_->getBindingDimensions(i);

    size_t volume = 1;
    for (int j = 0; j < dim.nbDims; ++j) volume *= dim.d[j];
    DataType type = engine_->getBindingDataType(i);
    switch (type) {
      case DataType::kFLOAT:
        volume *= sizeof(float);
        break;
      default:
        SPDLOG_ERROR("[TRTDetector] Do not support input type: {}", type);
        break;
    }

    void *device_memory = nullptr;
    cudaMalloc(&device_memory, volume);
    SPDLOG_WARN("cudaMalloc bingings_size_:{}, volume:{}", i, volume);

    bindings_.push_back(device_memory);
    bingings_size_.push_back(volume);

    SPDLOG_INFO("[TRTDetector] Binding {} : {}", i, engine_->getBindingName(i));
  }
  cudaStreamCreate(&stream_);

  return true;
}

TRTDetector::TRTDetector() {
  conf_thresh_ = Yolo::CONF_THRESH;
  nms_thresh_ = Yolo::NMS_THRESH;
  SPDLOG_TRACE("Constructed.");
}

TRTDetector::TRTDetector(const std::string &engine_file_path, float conf_thresh,
                         float nms_thresh) {
  SetEnginePath(engine_file_path);
  conf_thresh_ = conf_thresh;
  nms_thresh_ = nms_thresh;

  SPDLOG_TRACE("[TRTDetector] Constructed.");
}

TRTDetector::~TRTDetector() {
  SPDLOG_DEBUG("[TRTDetector] Destructing.");

  for (auto it = bindings_.begin(); it != bindings_.end(); ++it) cudaFree(*it);
  cudaStreamDestroy(stream_);

  delete[] data_;
  SPDLOG_DEBUG("[TRTDetector] Destructed.");
}

bool TRTDetector::LoadEngine(UniquePtr<ICudaEngine> engine) {
  SPDLOG_DEBUG("[TRTDetector] LoadEngine.");

  engine_ = std::move(engine);

  if (!engine_) {
    SPDLOG_ERROR("[TRTDetector] LoadEngine Fail.");
    return false;
  }
  SPDLOG_DEBUG("[TRTDetector] LoadEngine OK.");
  Init();
  return true;
}

void TRTDetector::SetEnginePath(const std::string &engine_path) {
  engine_path_ = engine_path;

  std::vector<char> engine_bin;
  std::ifstream engine_file(engine_path_, std::ios::binary);

  if (engine_file.good()) {
    SPDLOG_WARN("[TRTDetector] File is good.");
    engine_file.seekg(0, engine_file.end);
    engine_bin.resize(engine_file.tellg());
    engine_file.seekg(0, engine_file.beg);
    engine_file.read(engine_bin.data(), engine_bin.size());
    engine_file.close();
  } else {
    SPDLOG_ERROR("[TRTDetector] LoadEngine Fail. Could not open file. :{}",
                 engine_path);
    return;
  }

  auto runtime = UniquePtr<IRuntime>(createInferRuntime(logger_));

  engine_ = UniquePtr<ICudaEngine>(
      runtime->deserializeCudaEngine(engine_bin.data(), engine_bin.size()));

  SPDLOG_DEBUG("Onnx and engine file has been loaded");
}

void TRTDetector::Init() {
  CreateContex();
  InitMemory();
}

bool TRTDetector::TestInfer(const std::string &file_name) {
  SPDLOG_DEBUG("[TRTDetector] TestInfer.");
  cv::Mat raw = cv::imread(file_name);

  auto dets = Infer(raw);

  SPDLOG_WARN("{} [size] {}", file_name, dets.size());
  if (dets.size() > 0)

    for (std::size_t i = 0; i < dets.size(); i++) {
      component::Draw(raw, dets[i], i);
      component::PrintInfo(dets[i]);
    }

  cv::imwrite(component::Replace(file_name, "result"), raw);

  SPDLOG_DEBUG("[TRTDetector] TestInfer done.");
  return true;
}

std::vector<TRT::Detection> TRTDetector::Infer(const cv::Mat &raw) {
  SPDLOG_DEBUG("[TRTDetector] Infer.");
  PreProcess(raw);
  std::vector<float> output(bingings_size_.at(idx_out_) / sizeof(float));

  cudaMemcpyAsync(bindings_.at(idx_in_), data_, bingings_size_.at(idx_in_),
                  cudaMemcpyHostToDevice, stream_);

  context_->enqueue(Yolo::BATCH_SIZE, bindings_.data(), stream_, nullptr);

  cudaMemcpyAsync(output.data(), bindings_.at(idx_out_),
                  bingings_size_.at(idx_out_), cudaMemcpyDeviceToHost, stream_);

  cudaStreamSynchronize(stream_);

  auto dets = PostProcess(output, raw.cols, raw.rows);
  NonMaxSuppression(dets, nms_thresh_);

  SPDLOG_DEBUG("[TRTDetector] Infered.");
  return dets;
}
